import time
import numpy as np
import pygame
import get_rgb
from get_rgb import SimpleNN
import colorsys
import torch
import pandas as pd
from DMER2_0.songs_predict import process_single_song
from DMER2_0.CSBSmodel import CSBSModel
import torchaudio

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def predict_va_from_audio(audio_path, model_path='DMER2_0/model_logs/CSBSmodel_epoch10_10.pth'):
    """
    输入音频路径，输出该音频的VA值预测
    :param audio_path: 音频文件路径
    :param model_path: 训练好的模型权重路径
    :return: VA值（tensor）
    """
    input_tensor = process_single_song(audio_path)
    model = CSBSModel(1, 32, 3, 128, 2).to(DEVICE)
    state_dict = torch.load(model_path, weights_only=True)
    model.load_state_dict(state_dict)
    model.eval()
    with torch.no_grad():
        output = model(input_tensor)
    return output

def main():
    # 音乐文件路径
    music_path = "DMER2_0/datasets/test/chorus/1000.mp3"

    # 用模型预测整首歌的VA序列
    va_tensor = predict_va_from_audio(music_path)
    va_seq = va_tensor.cpu().numpy()  # shape: (帧数, 2)

    volumes = get_volume_per_frame(music_path)

    # 初始化音乐播放
    pygame.mixer.init()
    pygame.mixer.music.load(music_path)
    pygame.mixer.music.play()

    for idx, va in enumerate(va_seq):
        arousal, valence = va
        X_test = np.array([[arousal, valence]])
        H_predictions = get_rgb.predict(X_test, get_rgb.model_H).flatten()
        S_predictions = get_rgb.predict(X_test, get_rgb.model_S).flatten()
        # 用音量控制B_constant
        B_constant = volumes[idx] if idx < len(volumes) else 0.3
        print(f"音量: {B_constant:.4f}")
        # 转换为RGB colo
        color = colorsys.hsv_to_rgb(H_predictions[0], S_predictions[0], B_constant)
        r, g, b = int(color[0] * 255), int(color[1] * 255), int(color[2] * 255)
        print(f"Arousal: {arousal:.4f}, Valence: {valence:.4f}, RGB: ({r}, {g}, {b})")
        colors = [(r, g, b)] * get_rgb.LED_COUNT
        cmd = get_rgb.create_display_command(
            get_rgb.GROUP_ADDRESS,
            get_rgb.DEVICE_ADDRESS,
            get_rgb.PORT,
            get_rgb.LED_TYPE,
            get_rgb.REPEAT,
            colors
        )
        get_rgb.send_command(cmd)
        time.sleep(0.5)
        if not pygame.mixer.music.get_busy():
            break

    pygame.mixer.music.stop()


def get_volume_per_frame(audio_path, frame_duration=0.5):
    waveform, sr = torchaudio.load(audio_path)
    # 转为单通道
    if waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)
    frame_length = int(sr * frame_duration)
    total_length = waveform.shape[1]
    # 只保留能整除的部分
    num_frames = total_length // frame_length
    waveform = waveform[:, :num_frames * frame_length]
    # 使用unfold批量切帧
    frames = waveform.unfold(1, frame_length, frame_length)  # shape: (1, num_frames, frame_length)
    # 计算每帧RMS
    rms = frames.pow(2).mean(dim=2).sqrt().squeeze(0)  # shape: (num_frames,)
    # 归一化到0~1
    max_vol = rms.max().item() if rms.numel() > 0 else 1
    volumes = (rms / max_vol).tolist()
    return volumes

if __name__ == "__main__":
    main()